import numpy as np
import random

class CNSGA2():
    #-----------BEGIN-----------#
    # TODO: Implement the parameter-less constraint handling technique in NSGA-II

    def __init__(self, pop_size=100, n_var=5, pc=0.9, pm=0.1, eta_c=20, eta_m=20):
        self.pop_size = pop_size
        self.n_var = n_var
        self.pc = pc        # crossover probability
        self.pm = pm        # mutation probability
        self.eta_c = eta_c  # SBX distribution index
        self.eta_m = eta_m  # PM distribution index
        self.mating_pool = []
        self.offspring_pool = []
        self.ranks = []
        self.crowding_distances = []

        assert 0 <= self.pc <= 1
        assert 0 <= self.pm <= 1
        assert self.pop_size % 2 == 0

    def initialize(self, prob):
        x = prob.lower + (prob.upper - prob.lower) * np.random.rand(
            self.pop_size, prob.n_var)
        return x

    def constrained_fast_nondominated_sort(self, f, cv):
        N = f.shape[0]
        domination_counts = np.zeros(N)   # be dominated
        dominated_solutions = [[] for _ in range(N)]   # dominate
        rank = np.zeros(N)
        fronts = [[]]

        for p in range(N):
            for q in range(N):
                if p == q:
                    continue
                if cv[p] == 0 and cv[q] > 0:    # feasible dominates infeasible
                    dominated_solutions[p].append(q)
                elif cv[p] > 0 and cv[q] == 0:
                    domination_counts[p] += 1
                elif cv[p] == 0 and cv[q] == 0:   # both feasible
                    if np.all(f[p] <= f[q]) and np.any(f[p] < f[q]):
                        # p is not worse than q in all objectives and better at least 1
                        dominated_solutions[p].append(q)
                    elif np.all(f[q] <= f[p]) and np.any(f[q] < f[p]):
                        domination_counts[p] += 1
                elif cv[p] > 0 and cv[q] > 0:
                    if cv[p] < cv[q]:
                        dominated_solutions[p].append(q)
                    elif cv[p] > cv[q]:
                        domination_counts[p] += 1

            if domination_counts[p] == 0:
                rank[p] = 0
                fronts[0].append(p)

        i = 0
        while len(fronts[i]) > 0:
            next_front = []
            for p in fronts[i]:
                for q in dominated_solutions[p]:
                    domination_counts[q] -= 1
                    if domination_counts[q] == 0:   # the next front
                        rank[q] = i + 1
                        next_front.append(q)
            i += 1
            fronts.append(next_front)

        fronts.pop()
        return fronts, rank

    def fitness_assignment(self, f, cv):
        self.fronts, self.ranks = self.constrained_fast_nondominated_sort(f, cv)
        self.crowding_distances = np.zeros(len(f))
        n_obj = f.shape[1]

        for front in self.fronts:   # calculate crowding distance for every fronts
            if len(front) <= 1:
                for i in front:
                    self.crowding_distances[i] = np.inf
                continue
            for obj in range(n_obj):
                sorted_front = sorted(front, key=lambda x: f[x][obj])
                self.crowding_distances[sorted_front[0]] = np.inf
                self.crowding_distances[sorted_front[-1]] = np.inf

                f_max = f[sorted_front[-1]][obj]
                f_min = f[sorted_front[0]][obj]

                if f_max > f_min:
                    for i in range(1, len(sorted_front) - 1):
                        self.crowding_distances[sorted_front[i]] += (
                            f[sorted_front[i + 1]][obj] -
                            f[sorted_front[i - 1]][obj]
                        ) / (f_max - f_min)

        # normalize crowding_distances to [0,1]
        if np.all(np.isfinite(self.crowding_distances)):
            max_cd = np.max(self.crowding_distances)
            if max_cd > 0:
                self.crowding_distances /= max_cd


    def tournament_selection(self, population_x):
        mating_pool = []
        for _ in range(self.pop_size):
            # choose 2 individuals at random
            a, b = np.random.randint(0, self.pop_size, 2)
            # compare rank and crowding_distance
            if (self.ranks[a] < self.ranks[b] or
                (self.ranks[a] == self.ranks[b] and
                 self.crowding_distances[a] > self.crowding_distances[b])):
                winner = population_x[a]
            else:
                winner = population_x[b]
            mating_pool.append(winner)
        return mating_pool

    def crossover(self, pc):
        offspring_pool = []
        # for randomness
        np.random.shuffle(self.mating_pool)
        # in pairs
        for i in range(0, self.pop_size, 2):
            parent1 = self.mating_pool[i]
            parent2 = self.mating_pool[min(i + 1, self.pop_size - 1)]
            # init children
            child1, child2 = parent1.copy(), parent2.copy()

            # SBX crossover
            if np.random.random() <= pc:
                for j in range(len(parent1)):
                    u = np.random.random()
                    if u <= 0.5:
                        beta = (2 * u) ** (1 / (self.eta_c + 1))
                    else:
                        beta = (1 / (2 * (1 - u))) ** (1 / (self.eta_c + 1))
                    child1[j] = 0.5 * ((1 + beta) * parent1[j] + (1 - beta) * parent2[j])
                    child2[j] = 0.5 * ((1 - beta) * parent1[j] + (1 + beta) * parent2[j])

            # add small change if child equals parent
            if np.allclose(child1, parent1) and np.allclose(child1, parent2):
                child1 += np.random.normal(0, 1e-6, size=child1.shape)
            if np.allclose(child2, parent1) and np.allclose(child2, parent2):
                child2 += np.random.normal(0, 1e-6, size=child2.shape)


            offspring_pool.extend([child1, child2])
        # keep only pop_size children
        offspring_pool = offspring_pool[:self.pop_size]
        for child in offspring_pool:
            np.clip(child, self.lower, self.upper, out=child)
        return offspring_pool

    def mutation(self, pm):
        for i, individual in enumerate(self.offspring_pool):
            for j in range(len(individual)):
                if np.random.random() <= pm:
                    r = np.random.random()
                    if r < 0.5:
                        delta = (2 * r) ** (1 / (self.eta_m + 1)) - 1
                    else:
                        delta = 1 - (2 * (1 - r)) ** (1 / (self.eta_m + 1))
                    # mutation
                    individual[j] += delta * (self.upper[j] - self.lower[j])
                    individual[j] = np.clip(individual[j], self.lower[j], self.upper[j])

                    # move a bit inside
                    if abs(individual[j] - self.lower[j]) < 1e-6:
                        individual[j] += 1e-3
                    elif abs(individual[j] - self.upper[j]) < 1e-6:
                        individual[j] -= 1e-3

        return self.offspring_pool

    def environmental_selection(self, combined_population_list):
        population_x, offspring_x = combined_population_list
        population_x = np.array(population_x)
        offspring_x = np.array(offspring_x)
        combined_x = np.vstack([population_x, offspring_x])
        f, c = self.prob.evaluate(combined_x)   # define in MW7()
        cv = np.sum(c, axis=1)

        fronts, ranks = self.constrained_fast_nondominated_sort(f, cv)
        crowding_distances = np.zeros(len(combined_x))

        for front in fronts:   # calculate crowding distance again
            if len(front) <= 1:
                for i in front:
                    crowding_distances[i] = np.inf
                continue
            for obj in range(f.shape[1]):
                sorted_front = sorted(front, key=lambda x: f[x][obj])
                crowding_distances[sorted_front[0]] = np.inf
                crowding_distances[sorted_front[-1]] = np.inf
                f_max = f[sorted_front[-1]][obj]
                f_min = f[sorted_front[0]][obj]
                if f_max > f_min:
                    for i in range(1, len(sorted_front) - 1):
                        crowding_distances[sorted_front[i]] += (
                            f[sorted_front[i + 1]][obj] -
                            f[sorted_front[i - 1]][obj]
                        ) / (f_max - f_min)

        # sort by ranks and crowding_distances
        indices = np.lexsort((-crowding_distances, ranks))

        # survivor selection
        selected_indices = indices[:self.pop_size]
        selected_population = combined_x[selected_indices]
        selected_fitness = f[selected_indices]

        return [selected_population, selected_fitness]

    def run(self, prob, max_gen=250, seed=20250408):
        # send random seed
        if seed is not None:
            random.seed(seed)
            np.random.seed(seed)

        # init
        self.prob = prob
        self.lower = np.full(self.n_var, prob.lower)
        self.upper = np.full(self.n_var, prob.upper)

        population_x = self.initialize(prob)
        f, c = prob.evaluate(population_x)
        cv = np.sum(c, axis=1)

        # evolution loop
        for gen in range(1, max_gen):
            self.fitness_assignment(f, cv)
            self.mating_pool = self.tournament_selection(population_x)
            self.offspring_pool = self.crossover(self.pc)
            self.offspring_pool = self.mutation(self.pm)

            # evaluate offspring
            offspring_x = np.array(self.offspring_pool)
            offspring_f, offspring_c = prob.evaluate(offspring_x)
            offspring_cv = np.sum(offspring_c, axis=1)

            population_x, f = self.environmental_selection([population_x, offspring_x])
            c = prob.evaluate(population_x)[1]
            cv = np.sum(c, axis=1)

            # adaptive mutation based on population diversity
            diversity = np.std(population_x)
            if diversity < 1e-4:
                print("Low diversity: reduce eta_m")
                self.eta_m = max(5, self.eta_m*0.99)
            else:
                pass

            # print progress per 10 gens
            if gen % 10 == 0:
                print(f"Generation {gen}/{max_gen}")

        return population_x, f



    #-----------END-----------#